#!/bin/bash

# Script to restart all gollect trode instances
# Usage: ./restart_trodes.sh [num_trodes] [start_port] [config_dir]

# Default values
NUM_TRODES=${1:-3}
START_PORT=${2:-50051}
CONFIG_DIR=${3:-"./trode"}

# PID file that contains trode PIDs
PID_FILE="trode_pids.txt"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

echo -e "${BLUE}Restarting gollect trode instances...${NC}"
echo -e "${YELLOW}Configuration: $NUM_TRODES trodes starting on port $START_PORT${NC}"

# ============================================================================
# STEP 1: STOP EXISTING TRODES
# ============================================================================

echo -e "${BLUE}Step 1: Stopping existing trode instances...${NC}"

# Check if PID file exists
if [ ! -f "$PID_FILE" ]; then
    echo -e "${YELLOW}No PID file found. Looking for trode processes...${NC}"
    
    # Try to find and kill trode processes
    pids=$(pgrep -f "trode")
    if [ -z "$pids" ]; then
        echo -e "${GREEN}No trode processes found.${NC}"
    else
        echo -e "${YELLOW}Found trode processes: $pids${NC}"
        for pid in $pids; do
            echo -e "${YELLOW}Killing process $pid...${NC}"
            kill $pid 2>/dev/null
            if [ $? -eq 0 ]; then
                echo -e "${GREEN}Successfully killed process $pid${NC}"
            else
                echo -e "${RED}Failed to kill process $pid${NC}"
            fi
        done
    fi
else
    # Read PIDs from file and kill them
    if [ ! -s "$PID_FILE" ]; then
        echo -e "${YELLOW}PID file is empty. No trodes to kill.${NC}"
    else
        echo -e "${YELLOW}Reading PIDs from $PID_FILE...${NC}"
        
        killed_count=0
        while IFS= read -r pid; do
            if [ -n "$pid" ]; then
                echo -e "${YELLOW}Killing trode process $pid...${NC}"
                
                # Check if process still exists
                if kill -0 "$pid" 2>/dev/null; then
                    kill "$pid"
                    if [ $? -eq 0 ]; then
                        echo -e "${GREEN}Successfully killed process $pid${NC}"
                        ((killed_count++))
                    else
                        echo -e "${RED}Failed to kill process $pid${NC}"
                    fi
                else
                    echo -e "${YELLOW}Process $pid is no longer running${NC}"
                fi
            fi
        done < "$PID_FILE"
        
        echo -e "${GREEN}Killed $killed_count trode processes${NC}"
    fi
fi

# Clean up PID file
if [ -f "$PID_FILE" ]; then
    rm "$PID_FILE"
    echo -e "${BLUE}Removed PID file: $PID_FILE${NC}"
fi

# Clean up config files
echo -e "${BLUE}Cleaning up temporary config files...${NC}"
rm -f trode/config_*.yaml 2>/dev/null

# Clean up log files
echo -e "${BLUE}Cleaning up log files...${NC}"
rm -f trode/trode_*.log 2>/dev/null

echo -e "${GREEN}All existing trodes stopped and cleaned up!${NC}"

# ============================================================================
# STEP 2: WAIT FOR CLEANUP
# ============================================================================

echo -e "${BLUE}Step 2: Waiting for cleanup to complete...${NC}"
sleep 2

# ============================================================================
# STEP 3: START NEW TRODES
# ============================================================================

echo -e "${BLUE}Step 3: Starting new trode instances...${NC}"

# Check if trode binary exists
if [ ! -f "$CONFIG_DIR/trode" ]; then
    echo -e "${RED}Error: trode binary not found in $CONFIG_DIR${NC}"
    echo "Please build the trode first: cd $CONFIG_DIR && go build -o trode ."
    exit 1
fi

# Clear PID file
> "$PID_FILE"

# Function to create config for a specific port
create_config() {
    local port=$1
    local config_file="$CONFIG_DIR/config_${port}.yaml"
    
    cat > "$config_file" << EOF
cpu_interval: 5
memory_interval: 5
disk_interval: 5
io_interval: 5
process_interval: 5
user_interval: 5
port: $port
EOF
    echo "$config_file"
}

# Start trodes
for i in $(seq 1 $NUM_TRODES); do
    port=$((START_PORT + i - 1))
    
    echo -e "${YELLOW}Starting trode $i on port $port...${NC}"
    
    # Create config file for this port
    config_file=$(create_config $port)
    
    # Start trode in background
    cd "$CONFIG_DIR"
    ./trode -config "config_${port}.yaml" -port $port > "trode_${port}.log" 2>&1 &
    trode_pid=$!
    cd - > /dev/null
    
    # Save PID
    echo "$trode_pid" >> "$PID_FILE"
    
    echo -e "${GREEN}Trode $i started with PID $trode_pid on port $port${NC}"
    
    # Small delay to avoid port conflicts
    sleep 1
done

echo -e "${GREEN}All $NUM_TRODES trodes started successfully!${NC}"
echo -e "${BLUE}Trode PIDs saved to: $PID_FILE${NC}"

# Wait a bit for trodes to fully start up
echo -e "${YELLOW}Waiting for trodes to initialize...${NC}"
sleep 3

# Check if trodes are listening on their ports
echo -e "${YELLOW}Verifying trodes are listening on ports...${NC}"
for i in $(seq 1 $NUM_TRODES); do
    port=$((START_PORT + i - 1))
    if netstat -tuln 2>/dev/null | grep -q ":$port "; then
        echo -e "${GREEN}✓ Trode on port $port is listening${NC}"
    else
        echo -e "${YELLOW}⚠ Trode on port $port may not be ready yet${NC}"
    fi
done

echo -e "${GREEN}Trode restart completed successfully!${NC}"
echo -e "${BLUE}To monitor these trodes, run:${NC}"
echo -e "${YELLOW}cd monitor && ./gollect-monitor -addresses \"localhost:$START_PORT,localhost:$((START_PORT+1)),localhost:$((START_PORT+2))\" -interval 2s${NC}"
echo -e "${BLUE}To stop all trodes, run: ./kill_trodes.sh${NC}" 